{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练\n",
    "https://github.com/qqwweee/keras-yolo3\n",
    "\n",
    "训练所需框架配置：\n",
    "\n",
    " - Python 3.5.2\n",
    " - Keras 2.1.5\n",
    " - tensorflow 1.6.0\n",
    "\n",
    "\n",
    "### Training\n",
    "Generate your own annotation file and class names file.\n",
    "One row for one image;\n",
    "Row format: image_file_path box1 box2 ... boxN;\n",
    "Box format: x_min,y_min,x_max,y_max,class_id (no space).\n",
    "For VOC dataset, try python voc_annotation.py\n",
    "Here is an example:\n",
    "\n",
    "```\n",
    "path/to/img1.jpg 50,100,150,200,0 30,50,200,120,3\n",
    "path/to/img2.jpg 120,300,250,600,2\n",
    "...\n",
    "```\n",
    "\n",
    "Make sure you have run `python convert.py -w yolov3.cfg yolov3.weights model_data/yolo_weights.h5`\n",
    "The file model_data/yolo_weights.h5 is used to load pretrained weights.\n",
    "\n",
    "Modify train.py and start training.\n",
    "`python train.py`\n",
    "Use your trained weights or checkpoint weights with command line option --model model_file when using `yolo_video.py` Remember to modify `class path` or `anchor path`, with --classes class_file and --anchors anchor_file.\n",
    "\n",
    "训练经过：\n",
    "- 1 获得标注集\n",
    "- 2 yolov3.weight模型准备好\n",
    "- 3 修改class_file / anchor_file 的地址\n",
    "\n",
    "\n",
    "If you want to use original pretrained weights for YOLOv3:\n",
    "1. wget https://pjreddie.com/media/files/darknet53.conv.74\n",
    "2. rename it as darknet53.weights\n",
    "3. python convert.py -w darknet53.cfg darknet53.weights model_data/darknet53_weights.h5\n",
    "4. use model_data/darknet53_weights.h5 in train.py\n",
    "\n",
    "\n",
    "算法中的两步训练：\n",
    "\n",
    "- 1 先冻结一部分层数，迁移训练\n",
    "\n",
    "    Freeze the first 249 layers of total 252 layers.\n",
    "    \n",
    "    该步训练的结果本就不差\n",
    "    \n",
    "    \n",
    "- 2 全部解冻，再训练\n",
    "\n",
    "    该步骤观察训练结果，如果任何不好,可以继续训练\n",
    "    \n",
    "    Train longer if the result is not good.\n",
    "\n",
    "\n",
    "训练时间：\n",
    "\n",
    "第一次训练错误，xmin,ymin,xmax,ymax搞错了\n",
    "训练时长：35h\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.0.8'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import keras\n",
    "keras.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!python convert.py yolov3.cfg yolov3.weights model_data/yolo.h5   # 预测的时候用\n",
    "!python convert.py -w yolov3.cfg yolov3.weights model_data/yolo_weights.h5  #  训练的时候用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 自增长\n",
    "import tensorflow as tf\n",
    "from keras import backend as K\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth=True\n",
    "sess = tf.Session(config=config)\n",
    "K.set_session(sess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Create YOLOv3 model with 9 anchors and 33 classes.\n",
      "Load weights logs/003/ep009-loss33.297-val_loss32.851.h5.\n",
      "Freeze the first 249 layers of total 252 layers.\n",
      "Unfreeze all of the layers.\n",
      "Train on 44999 samples, val on 4999 samples, with batch size 16.\n",
      "Epoch 51/100\n",
      "2812/2812 [==============================] - 8630s 3s/step - loss: 29.1259 - val_loss: 27.4327\n",
      "Epoch 52/100\n",
      "2485/2812 [=========================>....] - ETA: 24:15 - loss: 27.0266"
     ]
    }
   ],
   "source": [
    "\n",
    "import numpy as np\n",
    "import keras.backend as K\n",
    "from keras.layers import Input, Lambda\n",
    "from keras.models import Model\n",
    "from keras.optimizers import Adam\n",
    "from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping\n",
    "\n",
    "from yolo3.model import preprocess_true_boxes, yolo_body, tiny_yolo_body, yolo_loss\n",
    "from yolo3.utils import get_random_data\n",
    "\n",
    "\n",
    "def _main(yolo_args):\n",
    "    annotation_path = yolo_args.annotation_path\n",
    "    log_dir = yolo_args.log_dir\n",
    "    classes_path = yolo_args.classes_path\n",
    "    anchors_path = yolo_args.anchors_path\n",
    "    class_names = get_classes(classes_path)\n",
    "    num_classes = len(class_names)\n",
    "    anchors = get_anchors(anchors_path)\n",
    "\n",
    "    input_shape = yolo_args.input_shape # multiple of 32, hw\n",
    "\n",
    "    is_tiny_version = len(anchors)==6 # default setting\n",
    "    if is_tiny_version:\n",
    "        model = create_tiny_model(input_shape, anchors, num_classes,\n",
    "            freeze_body=2, weights_path=yolo_args.weights_path) # 'model_data/tiny_yolo_weights.h5'\n",
    "    else:\n",
    "        model = create_model(input_shape, anchors, num_classes,\n",
    "            freeze_body=2, weights_path=yolo_args.weights_path) # make sure you know what you freeze 'model_data/yolo_weights.h5'\n",
    "\n",
    "    logging = TensorBoard(log_dir=log_dir)\n",
    "    checkpoint = ModelCheckpoint(log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',\n",
    "        monitor='val_loss', save_weights_only=True, save_best_only=True, period=3)\n",
    "    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1)\n",
    "    early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1)\n",
    "\n",
    "    val_split = yolo_args.val_split\n",
    "    with open(annotation_path) as f:\n",
    "        lines = f.readlines()\n",
    "    np.random.seed(10101)\n",
    "    np.random.shuffle(lines)\n",
    "    np.random.seed(None)\n",
    "    num_val = int(len(lines)*val_split)\n",
    "    num_train = len(lines) - num_val\n",
    "    '''\n",
    "    # 第一轮\n",
    "    '''\n",
    "    \n",
    "    # Train with frozen layers first, to get a stable loss.\n",
    "    # Adjust num epochs to your dataset. This step is enough to obtain a not bad model.\n",
    "    if yolo_args.stage_1_train:\n",
    "        model.compile(optimizer=Adam(lr=1e-3), loss={\n",
    "            # use custom yolo_loss Lambda layer.\n",
    "            'yolo_loss': lambda y_true, y_pred: y_pred})\n",
    "\n",
    "        batch_size = yolo_args.batch_size\n",
    "        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))\n",
    "        model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),\n",
    "                steps_per_epoch=max(1, num_train//batch_size),\n",
    "                validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),\n",
    "                validation_steps=max(1, num_val//batch_size),\n",
    "                epochs=yolo_args.epochs_stage_1,\n",
    "                initial_epoch=0,\n",
    "                callbacks=[logging, checkpoint])\n",
    "        model.save_weights(log_dir + 'trained_weights_stage_1.h5')\n",
    "    '''\n",
    "    # 第二轮\n",
    "    '''\n",
    "    # Unfreeze and continue training, to fine-tune.\n",
    "    # Train longer if the result is not good.\n",
    "    if yolo_args.finally_train:\n",
    "        for i in range(len(model.layers)):\n",
    "            model.layers[i].trainable = True\n",
    "        model.compile(optimizer=Adam(lr=1e-4), loss={'yolo_loss': lambda y_true, y_pred: y_pred}) # recompile to apply the change\n",
    "        print('Unfreeze all of the layers.')\n",
    "\n",
    "        batch_size = yolo_args.batch_size # note that more GPU memory is required after unfreezing the body\n",
    "        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))\n",
    "        model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),\n",
    "            steps_per_epoch=max(1, num_train//batch_size),\n",
    "            validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),\n",
    "            validation_steps=max(1, num_val//batch_size),\n",
    "            epochs=yolo_args.epochs_finally,  # 100\n",
    "            initial_epoch=50,\n",
    "            callbacks=[logging, checkpoint, reduce_lr, early_stopping])\n",
    "        model.save_weights(log_dir + 'trained_weights_final.h5')\n",
    "\n",
    "    # Further training if needed.\n",
    "\n",
    "\n",
    "def get_classes(classes_path):\n",
    "    '''loads the classes'''\n",
    "    with open(classes_path) as f:\n",
    "        class_names = f.readlines()\n",
    "    class_names = [c.strip() for c in class_names]\n",
    "    return class_names\n",
    "\n",
    "def get_anchors(anchors_path):\n",
    "    '''loads the anchors from a file'''\n",
    "    with open(anchors_path) as f:\n",
    "        anchors = f.readline()\n",
    "    anchors = [float(x) for x in anchors.split(',')]\n",
    "    return np.array(anchors).reshape(-1, 2)\n",
    "\n",
    "\n",
    "def create_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,\n",
    "            weights_path='model_data/yolo_weights.h5'):\n",
    "    '''create the training model'''\n",
    "    K.clear_session() # get a new session\n",
    "    image_input = Input(shape=(None, None, 3))\n",
    "    h, w = input_shape\n",
    "    num_anchors = len(anchors)\n",
    "\n",
    "    y_true = [Input(shape=(h//{0:32, 1:16, 2:8}[l], w//{0:32, 1:16, 2:8}[l], \\\n",
    "        num_anchors//3, num_classes+5)) for l in range(3)]\n",
    "\n",
    "    model_body = yolo_body(image_input, num_anchors//3, num_classes)\n",
    "    print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))\n",
    "\n",
    "    if load_pretrained:\n",
    "        model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)\n",
    "        print('Load weights {}.'.format(weights_path))\n",
    "        if freeze_body in [1, 2]:\n",
    "            # Freeze darknet53 body or freeze all but 3 output layers.\n",
    "            num = (185, len(model_body.layers)-3)[freeze_body-1]\n",
    "            for i in range(num): model_body.layers[i].trainable = False\n",
    "            print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))\n",
    "\n",
    "    model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',\n",
    "        arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5})(\n",
    "        [*model_body.output, *y_true])\n",
    "    model = Model([model_body.input, *y_true], model_loss)\n",
    "\n",
    "    return model\n",
    "\n",
    "def create_tiny_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,\n",
    "            weights_path='model_data/tiny_yolo_weights.h5'):\n",
    "    '''create the training model, for Tiny YOLOv3'''\n",
    "    K.clear_session() # get a new session\n",
    "    image_input = Input(shape=(None, None, 3))\n",
    "    h, w = input_shape\n",
    "    num_anchors = len(anchors)\n",
    "\n",
    "    y_true = [Input(shape=(h//{0:32, 1:16}[l], w//{0:32, 1:16}[l], \\\n",
    "        num_anchors//2, num_classes+5)) for l in range(2)]\n",
    "\n",
    "    model_body = tiny_yolo_body(image_input, num_anchors//2, num_classes)\n",
    "    print('Create Tiny YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))\n",
    "\n",
    "    if load_pretrained:\n",
    "        model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)\n",
    "        print('Load weights {}.'.format(weights_path))\n",
    "        if freeze_body in [1, 2]:\n",
    "            # Freeze the darknet body or freeze all but 2 output layers.\n",
    "            num = (20, len(model_body.layers)-2)[freeze_body-1]\n",
    "            for i in range(num): model_body.layers[i].trainable = False\n",
    "            print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))\n",
    "\n",
    "    model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',\n",
    "        arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.7})(\n",
    "        [*model_body.output, *y_true])\n",
    "    model = Model([model_body.input, *y_true], model_loss)\n",
    "\n",
    "    return model\n",
    "\n",
    "def data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes):\n",
    "    '''data generator for fit_generator'''\n",
    "    n = len(annotation_lines)\n",
    "    i = 0\n",
    "    while True:\n",
    "        image_data = []\n",
    "        box_data = []\n",
    "        for b in range(batch_size):\n",
    "            if i==0:\n",
    "                np.random.shuffle(annotation_lines)\n",
    "            image, box = get_random_data(annotation_lines[i], input_shape, random=True)\n",
    "            image_data.append(image)\n",
    "            box_data.append(box)\n",
    "            i = (i+1) % n\n",
    "        image_data = np.array(image_data)\n",
    "        box_data = np.array(box_data)\n",
    "        y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)\n",
    "        yield [image_data, *y_true], np.zeros(batch_size)\n",
    "\n",
    "def data_generator_wrapper(annotation_lines, batch_size, input_shape, anchors, num_classes):\n",
    "    n = len(annotation_lines)\n",
    "    if n==0 or batch_size<=0: return None\n",
    "    return data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes)\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    class yolo_args:\n",
    "        annotation_path = 'train.txt'\n",
    "        log_dir = 'logs/003/'\n",
    "        classes_path = 'model_data/class_file_en.txt'\n",
    "        anchors_path = 'model_data/yolo_anchors.txt'\n",
    "        input_shape = (416,416) # multiple of 32, hw\n",
    "        # 608*608  416*416  320*320\n",
    "        val_split = 0.1\n",
    "        batch_size = 16\n",
    "        epochs_stage_1 = 1\n",
    "        epochs_finally = 100\n",
    "        weights_path =   'logs/003/ep009-loss33.297-val_loss32.851.h5'# 可以使用'model_data/tiny_yolo_weights.h5' 也可以使用tiny_yolo的：'model_data/yolo_weights.h5'\n",
    "        stage_1_train = False\n",
    "        finally_train = True\n",
    "        \n",
    "        \n",
    "    # train\n",
    "    _main(yolo_args)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
